package main

import (
	"time"

	"xagent/internal/actors"
	"xagent/internal/knowledge"
	"xagent/internal/team"
	proto "xagent/proto/v1"

	"github.com/asynkron/protoactor-go/actor"
	"github.com/sirupsen/logrus"
)

func init() {
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp:   true,
		TimestampFormat: "2006-01-02 15:04:05",
	})
	logrus.SetLevel(logrus.DebugLevel)
}

func main() {
	logrus.Info("Starting novel writing example")

	// 创建 actor system
	system := actor.NewActorSystem()
	logrus.Info("Actor system created")

	// 创建知识库引擎
	rag := knowledge.NewRAGEngine(&knowledge.Config{
		EmbeddingModel: "text-embedding-ada-002",
		VectorStore:    "inmemory",
	})
	logrus.Info("Knowledge base engine created")

	// 创建写作团队
	novelTeam := team.NewTeam("novel-team", system, rag)
	logrus.WithField("team_name", "novel-team").Info("Novel writing team created")

	// 创建并雇佣各个角色
	roles := []struct {
		name        string
		role        string
		description string
	}{
		{"story-planner-1", "story_planner", "负责故事大纲和情节规划"},
		{"character-designer-1", "character_designer", "负责人物设计和性格塑造"},
		{"world-builder-1", "world_builder", "负责世界观和背景设定"},
		{"writer-1", "writer", "负责具体章节的写作"},
		{"editor-1", "editor", "负责文字润色和内容审查"},
		{"reviewer-1", "reviewer", "负责质量把控和反馈"},
	}

	for _, r := range roles {
		agent := actors.NewAIAgent(r.name, r.description)
		agent.SetKnowledgeBase(rag)
		if err := novelTeam.Hire(agent); err != nil {
			logrus.WithError(err).WithField("role", r.role).Error("Failed to hire agent")
			return
		}
		logrus.WithFields(logrus.Fields{
			"name": r.name,
			"role": r.role,
		}).Info("Agent hired")
	}

	// 定义写作任务
	tasks := []struct {
		id          string
		name        string
		description string
		assignAfter time.Duration
	}{
		{
			"plan-task-1",
			"CreateNovelOutline",
			"创建小说《重生在天际省》的整体大纲，包括主要情节、冲突和结局",
			0,
		},
		{
			"character-task-1",
			"DesignCharacters",
			"设计小说的主要角色，包括性格特征、背景故事和发展arc",
			2 * time.Second,
		},
		{
			"world-task-1",
			"BuildWorld",
			"构建天际省的世界观，包括魔法系统、社会结构和重要地点",
			4 * time.Second,
		},
		{
			"write-task-1",
			"WriteFirstChapter",
			"根据大纲写作第一章，重点描写主角重生的场景和初始冲突",
			6 * time.Second,
		},
		{
			"edit-task-1",
			"EditFirstChapter",
			"对第一章进行文字润色和内容优化，确保叙事流畅和细节丰富",
			8 * time.Second,
		},
		{
			"review-task-1",
			"ReviewFirstChapter",
			"评审第一章的质量，提供修改建议和改进方向",
			10 * time.Second,
		},
	}

	// 按顺序分配任务
	for _, t := range tasks {
		time.Sleep(t.assignAfter)

		task := &proto.Task{
			Id:          t.id,
			Name:        t.name,
			Description: t.description,
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Metadata: map[string]string{
				"created_at":     time.Now().Format(time.RFC3339),
				"novel_name":     "重生在天际省",
				"genre":          "奇幻",
				"target_readers": "游戏玩家，特别是上古卷轴系列的粉丝",
			},
		}

		logrus.WithFields(logrus.Fields{
			"task_id":   task.Id,
			"task_name": task.Name,
			"novel":     task.Metadata["novel_name"],
		}).Info("Assigning new task")

		if err := novelTeam.AssignTask(task); err != nil {
			logrus.WithError(err).WithField("task_id", task.Id).Error("Failed to assign task")
			continue
		}

		// 等待任务处理
		time.Sleep(2 * time.Second)

		// 获取并打印团队进度
		metrics := novelTeam.GetTeamMetrics()
		logrus.WithFields(logrus.Fields{
			"metrics": metrics,
		}).Info("Current team progress")
	}

	// 最终团队指标
	metrics := novelTeam.GetTeamMetrics()
	logrus.WithFields(logrus.Fields{
		"final_metrics": metrics,
	}).Info("Final team metrics")

	// 清理资源
	logrus.Info("Shutting down actor system")
	system.Shutdown()
}
